#!/usr/bin/env python3

from mlx_lm import load, generate # export TOKENIZERS_PARALLELISM=false
from subprocess import getoutput
from sys import argv

model, tokenizer = load("mlx-community/Meta-Llama-3-8B-Instruct-4bit")
# prompt = argv[1] # ./LLM_port_scan.py 'Отсканируй все порты на хосте 127.0.0.1'
# prompt = 'Отсканируй все порты на хосте 127.0.0.1'
# prompt = 'Какая версия ssh запущена на 127.0.0.1'
# prompt = 'Найти ssh алгоритмы доступные на 127.0.0.1'
prompt = input("Prompt: ")
messages = [ {"role": "user", 
"content": "Сгенерируй только команду для nmap без описания" + prompt} ]
input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
prompt = tokenizer.decode(input_ids)
response = generate(model, tokenizer, prompt=prompt).replace("`",'')
# print("LLM out1: ", response)
output = getoutput(response)
# print("Shell out1: ", output)
messages = [ {"role": "user", 
"content": "Кратко опиши по русски результаты:\n" + output} ]
input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
prompt = tokenizer.decode(input_ids)
print("LLM: ", end='')
for t in generate(model, tokenizer, prompt, max_tokens=512):
    print(t, end="", flush=True)
print()